# TODO: Add comment
# 
#base on these two command in R:
#image(y, 1:4, cbind(test,-test,test,-test), col = jet.colors(100), axes = FALSE,zlim=c(-200,200),xlab="UMR",ylab="")
#axis(2, 1:4,c("1","2","3","4"), las = 2, tick = FALSE, font=2) add axis and label at left of y-axis
#axis(4, 1:4,c("1","2","3","4"), las = 2, tick = FALSE, font=2) add axis and label at right side of y-axis
#axis(1, 100,c("1"), las = 1, tick = FALSE, font=2) add triangle at the center.. tick would add tick there..
# Author: yaping
# 2013-7-27
# 
#TO DO: use the right order of experiment list name
#

###############################################################################
library(colorRamps)
library(RColorBrewer)
jet.colors <-colorRampPalette(c("#00007F", "blue", "#007FFF", "cyan","#7FFF7F", "yellow", "#FF7F00", "red", "#7F0000"))

file_name_lists<-NULL
category_names<-NULL
sample_names<-NULL
experiment_names<-NULL
enrichScoreMax<-NULL
logscale=TRUE
tightMode=FALSE
capLimitPerc=-1
for (e in commandArgs(TRUE)) {
	ta = strsplit(e,"=",fixed=TRUE)
	if(! is.na(ta[[1]][2])) {
		if(ta[[1]][1] == "wd"){
			wd<-ta[[1]][2]  ## directory to make plot file
			
		}
		if(ta[[1]][1] == "prefix"){ ##prefix of the output file
			prefix<-ta[[1]][2]
		}
		if(ta[[1]][1] == "file_name_lists"){ ##each file name list should contain a list of experiment in this sample, aligned to one location; the second column should mentioned it is belong to "percentage" or "enrichment" system.
			file_name_lists<-c(file_name_lists,ta[[1]][2])
		}
		if(ta[[1]][1] == "category_names"){ ##number of different location bed files
			category_names<-c(category_names,ta[[1]][2])
		}
		if(ta[[1]][1] == "sample_names"){
			sample_names<-c(sample_names,ta[[1]][2])
		}
		if(ta[[1]][1] == "experiment_names"){
			experiment_names<-c(experiment_names,ta[[1]][2])
		}
		if(ta[[1]][1] == "step"){
			step<-as.numeric(ta[[1]][2])
		}
		if(ta[[1]][1] == "scale"){
			scale<-as.numeric(ta[[1]][2])
		}
		
		if(ta[[1]][1] == "bin_size_align"){
			bin_size_align<-as.numeric(ta[[1]][2])
		}
		if(ta[[1]][1] == "logscale"){
			logscale<-as.logical(ta[[1]][2])
		}
		if(ta[[1]][1] == "enrichScoreMax"){
			enrichScoreMax<-as.numeric(ta[[1]][2])
		}
		if(ta[[1]][1] == "capLimitPerc"){
			capLimitPerc<-as.numeric(ta[[1]][2])
		}
		if(ta[[1]][1] == "tightMode"){
			tightMode<-as.logical(ta[[1]][2])
		}

	}
}
setwd(wd)

valueLen<-as.integer((scale/bin_size_align)/(step/bin_size_align))*2+1

dataStore<-array(0, dim = c(length(category_names),length(sample_names),length(experiment_names),valueLen))

enrichScoreScaleMax=-Inf
enrichScoreScaleMin=Inf
expressionScoreScaleMax=-Inf
expressionScoreScaleMin=Inf
numericScoreScaleMax=-Inf
numericScoreScaleMin=Inf
percentageExpNum<-NULL
expressionExpNum<-NULL
numericExpNum<-NULL
enrichExpNum<-NULL
fileNum=1
for(file_name_list in file_name_lists){
	filesInfo<-read.table(file_name_list,sep="\t",header=F)
	for(fileNumInTable in 1:length(filesInfo[,1])){
		categoryNum=as.integer((fileNum-1)/(length(experiment_names)*length(sample_names)))+1
		sampleNum=as.integer((fileNum-1-(categoryNum-1)*length(experiment_names)*length(sample_names))/length(experiment_names))+1
		#fileNumInTable=ifelse(fileNum %% (length(experiment_names)) == 0 , length(experiment_names) ,fileNum %% (length(experiment_names)))
		valueGch<-NULL
		if(is.na(filesInfo[fileNumInTable,1]) | filesInfo[fileNumInTable,1] %in% ""){
			valueGch=rep(-Inf,valueLen) ##when there is no data, show it is fake no enrichment..
			
			
		}else if(filesInfo[fileNumInTable,2] %in% "numeric" & !is.na(as.numeric(filesInfo[fileNumInTable,1])) & is.numeric(as.numeric(filesInfo[fileNumInTable,1]))){ ##used for number of NDR or number of highly expressed genes
			if(filesInfo[fileNumInTable,1] %in% ""){
				valueGch=rep(-Inf,valueLen) 
			}
			valueGch<-rep(as.numeric(as.character(filesInfo[fileNumInTable,1])),valueLen)
			numericExpNum<-c(numericExpNum,fileNumInTable)
			#if(logscale){
			#	valueGch=ifelse(valueGch<=1,log2(1),log2(valueGch))
			#}
			if(!is.infinite(valueGch)){
				if(max(valueGch,na.rm=T)>numericScoreScaleMax){
					numericScoreScaleMax=max(valueGch,na.rm=T)
				}
				if(min(valueGch,na.rm=T)<numericScoreScaleMin){
					numericScoreScaleMin=min(valueGch,na.rm=T)
				}
			}
			
		}else{
			
			gch<-read.table(as.character(filesInfo[fileNumInTable,1]),sep="\t",header=F)		
			axisSeq<-seq(0-scale, scale, by=step)
			dataSeq<-seq(5+((length(gch[1,])-5)/2)-as.integer(scale/bin_size_align), 5+((length(gch[1,])-5)/2)+as.integer(scale/bin_size_align), by=as.integer(step/bin_size_align))
			if(capLimitPerc>0 & filesInfo[fileNumInTable,2] %in% "enrichement"){
				if(capLimitPerc>=1){
					stop("capLimitPerc should be between 0<capLimitPerc<1")
				}
				breaks=seq(0,1,capLimitPerc)
				capUplimit=quantile(unlist(gch[,5:length(gch[1,])]),probs=breaks,na.rm =T)[length(breaks)-1]
				gch[gch>capUplimit]=capUplimit
			}
			
			
			
			
			for(i in dataSeq){
				if(filesInfo[fileNumInTable,2] %in% "expression"){
					if(floor(i+step/(2*bin_size_align)-1)-ceiling(i-step/(2*bin_size_align)) <= 0){
						valueGch<-cbind(valueGch,median(gch[,ceiling(i-step/(2*bin_size_align))], na.rm=T))	
					}else{
						valueGch<-cbind(valueGch,median(gch[,ceiling(i-step/(2*bin_size_align)):floor(i+step/(2*bin_size_align)-1)], na.rm=T))	
					}
				}else{
					if(floor(i+step/(2*bin_size_align)-1)-ceiling(i-step/(2*bin_size_align)) <= 0){
						valueGch<-cbind(valueGch,mean(gch[,ceiling(i-step/(2*bin_size_align))], na.rm=T))	
					}else{
						valueGch<-cbind(valueGch,mean(colMeans(gch[,ceiling(i-step/(2*bin_size_align)):floor(i+step/(2*bin_size_align)-1)], na.rm=T), na.rm=T))	
					}
				}
				
				
			}
			if(filesInfo[fileNumInTable,2] %in% "enrichement"){ #not methylation percentage type of data
				#valueGch=ifelse(valueGch<=1,log2(1),log2(valueGch))
				if(logscale){
					valueGch=ifelse(valueGch<=1,log2(1),log2(valueGch))
				}
				if(!is.null(enrichScoreMax)){
					valueGch[valueGch>enrichScoreMax]=enrichScoreMax
				}
				
				
				if(max(valueGch,na.rm=T)>enrichScoreScaleMax){
					enrichScoreScaleMax=max(valueGch,na.rm=T)
				}
				if(min(valueGch,na.rm=T)<enrichScoreScaleMin){
					enrichScoreScaleMin=min(valueGch,na.rm=T)
				}
				enrichExpNum<-c(enrichExpNum,fileNumInTable)
			}else if(filesInfo[fileNumInTable,2] %in% "percentage"){
				percentageExpNum<-c(percentageExpNum,fileNumInTable)
			}else if(filesInfo[fileNumInTable,2] %in% "expression"){ #not methylation percentage type of data
				#valueGch=ifelse(valueGch<=1,log2(1),log2(valueGch))
				if(logscale){
					valueGch=ifelse(valueGch==0,log10(0.06),log10(valueGch))
				}
				
				if(max(valueGch,na.rm=T)>expressionScoreScaleMax){
					expressionScoreScaleMax=max(valueGch,na.rm=T)
				}
				if(min(valueGch,na.rm=T)<expressionScoreScaleMin){
					expressionScoreScaleMin=min(valueGch,na.rm=T)
				}
				expressionExpNum<-c(expressionExpNum,fileNumInTable)
			}
			
		}
		dataStore[categoryNum,sampleNum, fileNumInTable,]=valueGch
		fileNum=fileNum+1
	}
	
}

percentageExpNum<-percentageExpNum[!duplicated(percentageExpNum)]
expressionExpNum<-expressionExpNum[!duplicated(expressionExpNum)]
enrichExpNum<-enrichExpNum[!duplicated(enrichExpNum)]
numericExpNum<-numericExpNum[!duplicated(numericExpNum)]

#if(!is.null(enrichScoreMax)){
#	enrichScoreScaleMax=enrichScoreMax
#}

#style<-cbind(c(1,2),2+matrix(1:(length(sample_names)*2), 2, length(sample_names), byrow = TRUE))
percentageExpFrac=length(percentageExpNum)/length(experiment_names)
if(length(expressionExpNum)>0){
	expressionExpFrac=length(expressionExpNum)/length(experiment_names)
}else if(length(numericExpNum)>0){
	numericExpFrac=length(numericExpNum)/length(experiment_names)
}

titleFrac=percentageExpFrac/1.5
lht<-rep(c(titleFrac,percentageExpFrac,ifelse(length(expressionExpNum)>0,expressionExpFrac,numericExpFrac),ifelse(length(expressionExpNum)>0,1-titleFrac-percentageExpFrac-expressionExpFrac,1-titleFrac-percentageExpFrac-numericExpFrac)),3)
labelFrac=0.1
axisFarc=0.1
ncolsCategories=ceiling(length(category_names)/3)
lwd<-c(labelFrac,c(rep(c(rep((1-labelFrac-axisFarc)/length(sample_names),length(sample_names))),ncolsCategories-1),rep((1-labelFrac-axisFarc)/length(sample_names),length(sample_names)-1),(1-labelFrac-axisFarc)/length(sample_names)+axisFarc))


style<-matrix(1:((length(sample_names)*ncolsCategories+1)*4*3),  4*3, length(sample_names)*ncolsCategories+1, byrow = TRUE)
style[1:as.integer(length(style[,1])/3),1]=1
style[(as.integer(length(style[,1])/3)+1):as.integer(length(style[,1])*2/3),1]=2
style[(as.integer(length(style[,1])*2/3)+1):as.integer(length(style[,1])*3/3),1]=3

j=4
for(i in 2+(length(style[1,])-1)/ncolsCategories*(0:(ncolsCategories-1))){	
	for(r in seq(1,length(style[,1]),by=4)){

		for(z in i:(i+length(sample_names)-1)){	
			style[r,z]=j
		}
		j=j+1
		for(z in i:(i+length(sample_names)-1)){				
			style[(r+1),z]=j
			j=j+1
		}
		for(z in i:(i+length(sample_names)-1)){			
			style[(r+2),z]=j
			j=j+1
		}
		for(z in i:(i+length(sample_names)-1)){				
			style[(r+3),z]=j
			j=j+1
		}
	}
	
}


pdf(paste(prefix,"pdf",sep="."), paper="special", height=8*3, width=11*ncolsCategories*valueLen/501)
layout(style,heights = lht,widths=lwd)

par(mar=c(0.5, 4, 0.5, 1.5))
image(1, 0:100, matrix(0:100, 1, 101, byrow = TRUE), col = jet.colors(100), axes = FALSE,zlim=c(0,100),xlab="",ylab="Methylation/Accessibility",font.lab=2, cex.lab=0.75)
axis(2, seq(0,100,by=20), las = 2, tick = T, font=2, cex.axis=1.15)
if(length(expressionExpNum)>0){
	par(mar=c(0.5, 4, 0.5, 1.5))
	image(1, seq(expressionScoreScaleMin,expressionScoreScaleMax,by=(expressionScoreScaleMax-expressionScoreScaleMin)/100), matrix(seq(expressionScoreScaleMin,expressionScoreScaleMax,by=(expressionScoreScaleMax-expressionScoreScaleMin)/100), 1, 101, byrow = TRUE), col = jet.colors(100), axes = FALSE,zlim=c(expressionScoreScaleMin,expressionScoreScaleMax),xlab="",ylab="log10(RPKM)",font.lab=2, cex.lab=1.25)
	axis(2, seq(expressionScoreScaleMin,expressionScoreScaleMax,by=(expressionScoreScaleMax-expressionScoreScaleMin)/5), labels=format(seq(expressionScoreScaleMin,expressionScoreScaleMax,by=(expressionScoreScaleMax-expressionScoreScaleMin)/5),digits=1), las = 2, tick = T, font=2, cex.axis=1.15, cex.lab=1.5)
	
}
if(length(numericExpNum)>0){
	par(mar=c(0.5, 4, 0.5, 1.5))
	image(1, seq(numericScoreScaleMin,numericScoreScaleMax,by=(numericScoreScaleMax-numericScoreScaleMin)/100), matrix(seq(numericScoreScaleMin,numericScoreScaleMax,by=(numericScoreScaleMax-numericScoreScaleMin)/100), 1, 101, byrow = TRUE), col = jet.colors(100), axes = FALSE,zlim=c(numericScoreScaleMin,numericScoreScaleMax),xlab="",ylab="Percentage of genes",font.lab=2, cex.lab=1.25)
	axis(2, seq(numericScoreScaleMin,numericScoreScaleMax,by=(numericScoreScaleMax-numericScoreScaleMin)/5), labels=format(seq(numericScoreScaleMin,numericScoreScaleMax,by=(numericScoreScaleMax-numericScoreScaleMin)/5),digits=1), las = 2, tick = T, font=2, cex.axis=1.15, cex.lab=1.5)
	
}
par(mar=c(0.5, 4, 0.5, 1.5))
image(1, seq(enrichScoreScaleMin,enrichScoreScaleMax,by=(enrichScoreScaleMax-enrichScoreScaleMin)/100), matrix(seq(enrichScoreScaleMin,enrichScoreScaleMax,by=(enrichScoreScaleMax-enrichScoreScaleMin)/100), 1, 101, byrow = TRUE), col = jet.colors(100), axes = FALSE,zlim=c(enrichScoreScaleMin,enrichScoreScaleMax),xlab="",ylab="log2(Z-score)",font.lab=2, cex.lab=1.25)
axis(2, seq(enrichScoreScaleMin,enrichScoreScaleMax,by=(enrichScoreScaleMax-enrichScoreScaleMin)/5), labels=format(seq(enrichScoreScaleMin,enrichScoreScaleMax,by=(enrichScoreScaleMax-enrichScoreScaleMin)/5),digits=1), las = 2, tick = T, font=2, cex.axis=1.15, cex.lab=1.5)


for(categoryNum in 1:length(category_names)){
	
	
	plot.new()
	par(mar=c(0.5, 0, 0.5, 0.5))
	segments(-1,par("usr")[1],par("usr")[2],par("usr")[1],lwd=2)
	text((par("usr")[2]-par("usr")[1])/2,(par("usr")[4]-par("usr")[3])/2, category_names[categoryNum], font=2, cex=2)
	for(sampleNum in 1:(length(sample_names)-1)){
		#par(oma=c(0, 0, 0, 0),mar=c(0, 1, 2, 2))
		par(mar=c(0.5, 0.5, 2, 1))
		#image(1:valueLen, 1:length(percentageExpNum), apply(t(dataStore[categoryNum,sampleNum,percentageExpNum,]), 1, rev), col = jet.colors(100), axes = FALSE,zlim=c(0,100),xlab="",ylab="")
		image(1:valueLen, 1:length(percentageExpNum), t(dataStore[categoryNum,sampleNum,percentageExpNum,]), col = jet.colors(100), axes = FALSE,zlim=c(0,100),xlab="",ylab="")
		#axis(4, 1:length(percentageExpNum),experiment_names[percentageExpNum], las = 2, tick = FALSE, font=2)
		axis(1, valueLen/2, "", las = 2, tick = TRUE, font=2,cex.axis=2)
		title(sample_names[sampleNum], font=2, cex=1.5)
	}
	
	#image(1:valueLen, 1:length(percentageExpNum), apply(t(dataStore[categoryNum,length(sample_names),percentageExpNum,]), 1, rev), col = jet.colors(100), axes = FALSE,zlim=c(0,100),xlab="",ylab="")
	if(categoryNum > length(category_names)-3){
		par(mar=c(0.5, 0.5, 2, 8))
		
	}else{
		par(mar=c(0.5, 0.5, 2, 1))
	}
	image(1:valueLen, 1:length(percentageExpNum), t(dataStore[categoryNum,length(sample_names),percentageExpNum,]), col = jet.colors(100), axes = FALSE,zlim=c(0,100),xlab="",ylab="")
	if(categoryNum > length(category_names)-3){
		axis(4, 1:length(percentageExpNum),experiment_names[percentageExpNum], las = 2, tick = FALSE, font=2,cex.lab=1.5,cex.axis=1.15)
	}
	axis(1, valueLen/2, "", las = 2, tick = TRUE, font=2,cex.axis=2)
	title(sample_names[length(sample_names)])
	if(length(expressionExpNum)>0){
		for(sampleNum in 1:(length(sample_names)-1)){
			#par(oma=c(0, 0, 0, 0),mar=c(0, 1, 0, 2))
			par(mar=c(0.5, 0.5, 0, 1))
			#image(1:valueLen, 1:(length(experiment_names)-length(percentageExpNum)), apply(t(dataStore[categoryNum,sampleNum,-percentageExpNum,]), 1, rev), col = jet.colors(100), axes = FALSE,zlim=c(enrichScoreScaleMin,enrichScoreScaleMax),xlab="",ylab="")
			image(1:valueLen, 1:(length(expressionExpNum)), t(dataStore[categoryNum,sampleNum,expressionExpNum,]), col = jet.colors(100), axes = FALSE,zlim=c(expressionScoreScaleMin,expressionScoreScaleMax),xlab="",ylab="")
			#axis(4, 1:(length(experiment_names)-length(percentageExpNum)),experiment_names[-percentageExpNum], las = 2, tick = FALSE, font=2)
			axis(1, valueLen/2, "", las = 2, tick = TRUE, font=2,cex.axis=2)
		}
		
		#image(1:valueLen, 1:(length(experiment_names)-length(percentageExpNum)), apply(t(dataStore[categoryNum,length(sample_names),-percentageExpNum,]), 1, rev), col = jet.colors(100), axes = FALSE,zlim=c(enrichScoreScaleMin,enrichScoreScaleMax),xlab="",ylab="")
		
		if(categoryNum > length(category_names)-3){
			par(mar=c(0.5, 0.5, 0, 8))
			
		}else{
			par(mar=c(0.5, 0.5, 0, 1))
		}
		image(1:valueLen, 1:(length(expressionExpNum)), t(dataStore[categoryNum,length(sample_names),expressionExpNum,]), col = jet.colors(100), axes = FALSE,zlim=c(expressionScoreScaleMin,expressionScoreScaleMax),xlab="",ylab="")
		if(categoryNum > length(category_names)-3){
			axis(4, 1:(length(expressionExpNum)),experiment_names[expressionExpNum], las = 2, tick = FALSE, font=2,cex.axis=1.15)
		}
		axis(1, valueLen/2, "", las = 2, tick = TRUE, font=2,cex.axis=2)
	}
	
	if(length(numericExpNum)>0){
		for(sampleNum in 1:(length(sample_names)-1)){
			#par(oma=c(0, 0, 0, 0),mar=c(0, 1, 0, 2))
			par(mar=c(0.5, 0.5, 0, 1))
			#image(1:valueLen, 1:(length(experiment_names)-length(percentageExpNum)), apply(t(dataStore[categoryNum,sampleNum,-percentageExpNum,]), 1, rev), col = jet.colors(100), axes = FALSE,zlim=c(enrichScoreScaleMin,enrichScoreScaleMax),xlab="",ylab="")
			image(1:valueLen, 1:(length(numericExpNum)), t(dataStore[categoryNum,sampleNum,numericExpNum,]), col = jet.colors(100), axes = FALSE,zlim=c(numericScoreScaleMin,numericScoreScaleMax),xlab="",ylab="")
			#axis(4, 1:(length(experiment_names)-length(percentageExpNum)),experiment_names[-percentageExpNum], las = 2, tick = FALSE, font=2)
			axis(1, valueLen/2, "", las = 2, tick = TRUE, font=2,cex.axis=2)
		}
		
		#image(1:valueLen, 1:(length(experiment_names)-length(percentageExpNum)), apply(t(dataStore[categoryNum,length(sample_names),-percentageExpNum,]), 1, rev), col = jet.colors(100), axes = FALSE,zlim=c(enrichScoreScaleMin,enrichScoreScaleMax),xlab="",ylab="")

		if(categoryNum > length(category_names)-3){
			par(mar=c(0.5, 0.5, 0, 8))
			
		}else{
			par(mar=c(0.5, 0.5, 0, 1))
		}
		image(1:valueLen, 1:(length(numericExpNum)), t(dataStore[categoryNum,length(sample_names),numericExpNum,]), col = jet.colors(100), axes = FALSE,zlim=c(numericScoreScaleMin,numericScoreScaleMax),xlab="",ylab="")
		if(categoryNum > length(category_names)-3){
			axis(4, 1:(length(numericExpNum)),experiment_names[numericExpNum], las = 2, tick = FALSE, font=2,cex.axis=1.15)
		}
		axis(1, valueLen/2, "", las = 2, tick = TRUE, font=2,cex.axis=2)
	}
	
	
	for(sampleNum in 1:(length(sample_names)-1)){
		#par(oma=c(0, 0, 0, 0),mar=c(0, 1, 0, 2))
		par(mar=c(0.5, 0.5, 0, 1))
		#image(1:valueLen, 1:(length(experiment_names)-length(percentageExpNum)), apply(t(dataStore[categoryNum,sampleNum,-percentageExpNum,]), 1, rev), col = jet.colors(100), axes = FALSE,zlim=c(enrichScoreScaleMin,enrichScoreScaleMax),xlab="",ylab="")
		image(1:valueLen, 1:(length(enrichExpNum)), t(dataStore[categoryNum,sampleNum,enrichExpNum,]), col = jet.colors(100), axes = FALSE,zlim=c(enrichScoreScaleMin,enrichScoreScaleMax),xlab="",ylab="")
		#axis(4, 1:(length(experiment_names)-length(percentageExpNum)),experiment_names[-percentageExpNum], las = 2, tick = FALSE, font=2)
		axis(1, valueLen/2, "", las = 2, tick = TRUE, font=2,cex.axis=2)
	}
	
	#image(1:valueLen, 1:(length(experiment_names)-length(percentageExpNum)), apply(t(dataStore[categoryNum,length(sample_names),-percentageExpNum,]), 1, rev), col = jet.colors(100), axes = FALSE,zlim=c(enrichScoreScaleMin,enrichScoreScaleMax),xlab="",ylab="")
	if(categoryNum > length(category_names)-3){
		par(mar=c(0.5, 0.5, 0, 8))
		
	}else{
		par(mar=c(0.5, 0.5, 0, 1))
	}
	image(1:valueLen, 1:(length(enrichExpNum)), t(dataStore[categoryNum,length(sample_names),enrichExpNum,]), col = jet.colors(100), axes = FALSE,zlim=c(enrichScoreScaleMin,enrichScoreScaleMax),xlab="",ylab="")
	if(categoryNum > length(category_names)-3){
		axis(4, 1:(length(enrichExpNum)),experiment_names[enrichExpNum], las = 2, tick = FALSE, font=2,cex.axis=1.15)	
	}
	
	axis(1, valueLen/2, "", las = 2, tick = TRUE, font=2,cex.axis=2)
}
dev.off()


